import matplotlib.pyplot as plt
import numpy as np
import librosa.display
from sklearn.metrics import confusion_matrix
import seaborn as sns
from typing import List, Optional

## 可视化
# 混淆矩阵绘制（支持归一化）
# 训练过程曲线可视化
# 多类别频谱特征对比

def plot_confusion_matrix(
        y_true: np.ndarray,
        y_pred: np.ndarray,
        class_names: List[str],
        title: str = "Confusion Matrix",
        normalize: bool = True,
        save_path: Optional[str] = None
) -> None:
    """绘制混淆矩阵"""
    cm = confusion_matrix(y_true, y_pred)
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

    plt.figure(figsize=(10, 8))
    sns.heatmap(
        cm,
        annot=True,
        fmt=".2f" if normalize else "d",
        cmap="Blues",
        xticklabels=class_names,
        yticklabels=class_names
    )
    plt.title(title)
    plt.ylabel('True label')
    plt.xlabel('Predicted label')

    if save_path:
        plt.savefig(save_path, bbox_inches='tight')
    else:
        plt.show()


def plot_training_history(
        history: dict,
        metrics: List[str] = ['loss', 'accuracy'],
        save_path: Optional[str] = None
) -> None:
    """绘制训练历史曲线"""
    plt.figure(figsize=(12, 5))

    for i, metric in enumerate(metrics):
        plt.subplot(1, len(metrics), i + 1)
        plt.plot(history.get(f'train_{metric}', []), label=f'Train {metric}')
        plt.plot(history.get(f'val_{metric}', []), label=f'Validation {metric}')
        plt.title(f'Model {metric.capitalize()}')
        plt.ylabel(metric)
        plt.xlabel('epoch')
        plt.legend()

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, bbox_inches='tight')
    else:
        plt.show()


def plot_feature_comparison(
        features: List[np.ndarray],
        labels: List[str],
        sr: int = 44100,
        hop_length: int = 512
) -> None:
    """对比不同类别的频谱特征"""
    plt.figure(figsize=(15, 8))

    for i, (feature, label) in enumerate(zip(features, labels)):
        plt.subplot(2, len(features), i + 1)
        librosa.display.specshow(
            feature,
            sr=sr,
            hop_length=hop_length,
            x_axis='time',
            y_axis='mel'
        )
        plt.colorbar(format='%+2.0f dB')
        plt.title(label)

    plt.tight_layout()
    plt.show()

## 使用示例
# from utils.visualization import plot_confusion_matrix
#
# plot_confusion_matrix(
#     y_true=true_labels,
#     y_pred=pred_labels,
#     class_names=["excavator", "shovel", "drill"]
# )